import ext4.data as data
import ext4.jbd2 as journal


class SuperBlock(data.HexArea):
    rules = (
        (data.le32, 1, "inodes_count"),
        (data.le32, 1, "blocks_count_lo"),
        (data.le32, 1, "r_blocks_count_lo"),
        (data.le32, 1, "free_blocks_count_lo"),
        (data.le32, 1, "free_inodes_count"),
        (data.le32, 1, "first_data_block"),
        (data.le32, 1, "log_block_size"),
        (data.le32, 1, "log_cluster_size"),
        (data.le32, 1, "blocks_per_group"),
        (data.le32, 1, "clusters_per_group"),
        (data.le32, 1, "inodes_per_group"),
        (data.le32, 1, "mtime"),
        (data.le32, 1, "wtime"),
        (data.le16, 1, "mnt_count"),
        (data.le16, 1, "max_mnt_count"),
        (data.le16, 1, "magic"),
        (data.le16, 1, "state"),
        (data.le16, 1, "errors"),
        (data.le16, 1, "minor_rev_level"),
        (data.le32, 1, "lastcheck"),
        (data.le32, 1, "checkinterval"),
        (data.le32, 1, "creator_os"),
        (data.le32, 1, "rev_level"),
        (data.le16, 1, "def_resuid"),
        (data.le16, 1, "def_resgid"),
        (data.le32, 1, "first_ino"),
        (data.le16, 1, "inode_size"),
        (data.le16, 1, "block_group_nr"),
        (data.SuperBlock_Feature_Compat, 1, "feature_compat"),
        (data.SuperBlock_Feature_Incompat, 1, "feature_incompat"),
        (data.SuperBlock_Feature_RO_Compat, 1, "feature_ro_compat"),
        (data.le8, 16, "uuid"),
        (str, 16, "volume_name"),
        (str, 64, "last_mounted"),
        (data.le32, 1, "algorithm_usage_bitmap"),
        (data.le8, 1, "prealloc_blocks"),
        (data.le8, 1, "prealloc_dir_blocks"),
        (data.le16, 1, "reserved_gdt_blocks"),
        (data.le8, 16, "journal_uuid"),
        (data.le32, 1, "journal_inum"),
        (data.le32, 1, "journal_dev"),
        (data.le32, 1, "last_orphan"),
        (data.le32, 4, "hash_seed"),
        (data.le8, 1, "def_hash_version"),
        (data.le8, 1, "jnl_backup_type"),
        (data.le16, 1, "desc_size"),  # group size
        (data.le32, 1, "default_mount_opts"),
        (data.le32, 1, "first_meta_bg"),
        (data.le32, 1, "mkfs_time"),
        (data.le32, 17, "jnl_blocks"),
        (data.le32, 1, "blocks_count_hi"),
        (data.le32, 1, "r_blocks_count_hi"),
        (data.le32, 1, "free_inodes_count_hi"),
        (data.le16, 1, "min_extra_isize"),
        (data.le16, 1, "want_extra_isize"),
        (data.SuperBlock_Flags, 1, "flags"),
        (data.le16, 1, "raid_stride"),
        (data.le16, 1, "mmp_update_interval"),
        (data.le64, 1, "mmp_block"),
        (data.le32, 1, "raid_stripe_width"),
        (data.le8, 1, "log_groups_per_flex"),
        (data.SuperBlock_Checksum_Type, 1, "checksum_type"),
        (data.le8, 1, "encryption_level"),
        (data.le8, 1, "reserved_pad"),
        (data.le64, 1, "kbytes_written"),
        (data.le32, 1, "snapshot_inum"),
        (data.le32, 1, "snapshot_id"),
        (data.le64, 1, "snapshot_r_blocks_count"),
        (data.le32, 1, "snapshot_list"),
        (data.le32, 1, "error_count"),
        (data.le32, 1, "first_error_time"),
        (data.le32, 1, "first_error_ino"),
        (data.le64, 1, "first_error_block"),
        (str, 32, "first_error_func"),
        (data.le32, 1, "first_error_line"),
        (data.le32, 1, "last_error_time"),
        (data.le32, 1, "last_error_ino"),
        (data.le32, 1, "last_error_line"),
        (data.le64, 1, "last_error_block"),
        (str, 32, "last_error_func"),
        (data.le8, 64, "mount_opts"),
        (data.le32, 1, "usr_quota_inum"),
        (data.le32, 1, "grp_quota_inum"),
        (data.le32, 1, "overhead_clusters"),
        (data.le32, 2, "backup_bgs"),
        (data.le8, 4, "encrypt_algos"),
        (data.le8, 16, "encrypt_pw_salt"),
        (data.le32, 1, "lpf_ino"),
        (data.le32, 1, "prj_quota_inum"),
        (data.le32, 1, "checksum_seed"),
        (data.le8, 1, "wtime_hi"),
        (data.le8, 1, "mtime_hi"),
        (data.le8, 1, "mkfs_time_hi"),
        (data.le8, 1, "lastcheck_hi"),
        (data.le8, 1, "first_error_time_hi"),
        (data.le8, 1, "last_error_time_hi"),
        (data.le8, 1, "first_error_errcode"),
        (data.le8, 1, "last_error_errcode"),
        (data.le16, 1, "encoding"),
        (data.le16, 1, "encoding_flags"),
        (data.le32, 95, "reserved"),
        (data.le32, 1, "checksum"),
    )

    def __init__(self, bytes) -> None:
        super().__init__(bytes)

        self.blocks_count = data.le64(
            lo=self.blocks_count_lo, hi=self.blocks_count_hi
        ).value


class Group(data.HexArea):
    rules = (
        (data.le32, 1, "block_bitmap_lo"),
        (data.le32, 1, "inode_bitmap_lo"),
        (data.le32, 1, "inode_table_lo"),
        (data.le16, 1, "free_blocks_count_lo"),
        (data.le16, 1, "free_inodes_count_lo"),
        (data.le16, 1, "used_dirs_count_lo"),
        (data.le16, 1, "flags"),
        (data.le32, 1, "exclude_bitmap_lo"),
        (data.le16, 1, "block_bitmap_csum_lo"),
        (data.le16, 1, "inode_bitmap_csum_lo"),
        (data.le16, 1, "itable_unused_lo"),
        (data.le16, 1, "checksum"),
        (data.le32, 1, "block_bitmap_hi"),
        (data.le32, 1, "inode_bitmap_hi"),
        (data.le32, 1, "inode_table_hi"),
        (data.le16, 1, "free_blocks_count_hi"),
        (data.le16, 1, "free_inodes_count_hi"),
        (data.le16, 1, "used_dirs_count_hi"),
        (data.le16, 1, "itable_unused_hi"),
        (data.le32, 1, "exclude_bitmap_hi"),
        (data.le16, 1, "block_bitmap_csum_hi"),
        (data.le16, 1, "inode_bitmap_csum_hi"),
        (data.le32, 1, "reserved"),
    )

    extern_str = ("inode_bitmap", "block_bitmap", "inode_table", "blocks_count")

    def __init__(self, bytes) -> None:
        super().__init__(bytes)

        self.inode_bitmap = data.le64(lo=self.inode_bitmap_lo, hi=self.inode_bitmap_hi)
        self.block_bitmap = data.le64(lo=self.block_bitmap_lo, hi=self.block_bitmap_hi)

        self.inode_table = data.le64(lo=self.inode_table_lo, hi=self.inode_table_hi)
        self.blocks_count = data.le64(
            lo=self.free_blocks_count_lo, hi=self.free_blocks_count_hi
        )


class InodeFile(data.HexArea):
    rules = (
        (data.le32, 1, "inode"),
        (data.le16, 1, "rec_len"),
        (data.u8, 1, "name_len"),
        (data.FileType, 1, "file_type"),
    )

    extern_str = ("name",)

    def __init__(self, bytes, load_name=True) -> None:
        super().__init__(bytes)

        """ref: __ext4_check_dir_entry"""
        assert self.rec_len % 4 == 0

        if load_name is True:
            self.load_rules(((str, self.name_len, "name"),))

    def set_name(self, bytes):
        self.load_rules(((str, self.name_len, "name"),), bytes=bytes)

    def __str__(self) -> str:
        return "%7d %8s %s" % (self.inode.value, str(self.file_type), self.name)


class InodeXttr(data.HexArea):
    """ref: ext4_xattr_get"""

    rules = (
        (data.u8, 1, "name_len"),
        (data.u8, 1, "name_index"),
        (data.le16, 1, "value_offs"),
        (data.le32, 1, "value_inum"),
        (data.le32, 1, "value_size"),
        (data.le32, 1, "hash"),
    )

    extern_str = ("name",)

    def __init__(self, bytes, load_name=True, load_value=False) -> None:
        super().__init__(bytes)

        if load_name is True:
            self.load_rules(((str, self.name_len, "name"),))

        if load_value is True:
            self.value = bytes[
                self.value_offs.value : self.value_offs.value + self.value_size.value
            ]

    def set_value(self, bytes):
        """use another bytes, for load from block not inode"""
        self.value = bytes[
            self.value_offs.value : self.value_offs.value + self.value_size.value
        ]


class Inode(data.HexArea):
    rules = (
        (data.InodeMode, 1, "mode"),
        (data.le16, 1, "uid"),
        (data.le32, 1, "size_lo"),
        (data.Timestamp, 1, "atime"),
        (data.Timestamp, 1, "ctime"),
        (data.Timestamp, 1, "mtime"),
        (data.Timestamp, 1, "dtime"),
        (data.le16, 1, "gid"),
        (data.le16, 1, "links_count"),
        (data.le32, 1, "blocks_lo"),
        (data.InodeFlag, 1, "flags"),
        (data.Inode_Osd1, 1, "osd1"),
        (data.InodeBlock, 1, "block"),
        (data.le32, 1, "generation"),
        (data.le32, 1, "file_acl_lo"),
        (data.le32, 1, "size_high"),
        (data.le32, 1, "obso_faddr"),
        (data.Inode_Osd2, 1, "osd2"),
        (data.le16, 1, "extra_isize"),
        (data.le16, 1, "checksum_hi"),
        (data.Timestamp_Extra, 1, "ctime_extra"),
        (data.Timestamp_Extra, 1, "mtime_extra"),
        (data.Timestamp_Extra, 1, "atime_extra"),
        (data.Timestamp, 1, "crtime"),
        (data.Timestamp_Extra, 1, "crtime_extra"),
        (data.le32, 1, "version_hi"),
        (data.le32, 1, "projid"),
    )

    extern_str = ("checksum", "size", "blocks", "extents")

    INODE_0_NOT_EXIST = (0,)
    INODE_LIST_OF_DEFECTIVE_BLOCKS = (1,)
    INODE_ROOT_DIR = 2
    INODE_USER_QUOTA = 3
    INODE_GROUP_QUOTA = 4
    INODE_BOOT_LOADER = 5
    INODE_UNDELETE_DIR = 6
    INODE_RESERVED_GROUP_DESC = 7
    INODE_JOURNAL = 8
    INODE_EXCLUDE_FOR_SNAPSHOTS = 9
    INODE_REPLICA_FOR_NON_UPSTREAM_FEATURE = 10
    INODE_LOST_FOUND_DIR = 11

    def __init__(self, bytes, ext4) -> None:
        super().__init__(bytes)
        self.ext4 = ext4
        self.offset = 0  # offset of self.size
        self.offset_extents = 0  # offset of self.extents.extens

        # offset of block at self.extents.extens[extent_index]
        self.offset_extent_len = 0

        self.checksum = data.le32(lo=self.osd2.linux2.checksum_lo, hi=self.checksum_hi)
        self.size = data.le64(lo=self.size_lo, hi=self.size_high)
        self.blocks = data.le64(lo=self.blocks_lo, hi=self.osd2.linux2.blocks_high)
        self.uid = data.le32(lo=self.uid, hi=self.osd2.linux2.uid_high)
        self.gid = data.le32(lo=self.gid, hi=self.osd2.linux2.gid_high)
        self.file_acl = data.le64(
            lo=self.file_acl_lo, hi=self.osd2.linux2.file_acl_high
        )

        if self.flags.has("EXTENTS"):
            self.block.choose = "extent"
        elif self.mode.islnk() and self.size <= self.blocks.length:
            self.block.choose = "link"
            self.block.link.length = self.size.value

        self.extra_data = bytes[128 + self.extra_isize :]  # for xattr

        """init extents"""
        if self.flags.has("EXTENTS"):
            self._parse_extent_header(self.block.extent)
            if self.size > 0:
                assert len(self.extents) >= 1, "extents too small"
                self.extents.sort()

                offset = 0
                for i in range(len(self.extents) - 1):
                    offset += self.extents[i].length
                    assert offset == self.extents[i + 1].block


    def _parse_extent_header(self, header):
        self.extents = data.Array()

        if header.depth == 0 and header.entries > 0:
            for i in range(header.entries.value):
                self.extents.append(
                    data.InodeExtent(
                        header.bytes[
                            data.InodeExtentHeader.length
                            + data.InodeExtent.length
                            * i : data.InodeExtentHeader.length
                            + data.InodeExtent.length * (i + 1)
                        ]
                    )
                )
        elif header.depth > 0 and header.entries > 0:
            """not implemented"""
            for i in range(0, header.entries.value):
                idx = data.InodeExtentIdx(
                    header.bytes[
                        data.InodeExtentHeader.length
                        + data.InodeExtentIdx.length * i : data.InodeExtentHeader.length
                        + data.InodeExtentIdx.length * (i + 1)
                    ]
                )

                self.ext4.fd.seek(idx.leaf * self.ext4.blocksize)
                header = data.InodeExtentHeader(self.ext4.fd.read(self.ext4.blocksize))
                self._parse_extent_header(header)

    def read(self, size=None):
        if self.size == 0 or not hasattr(self, "extents"):
            return b""

        ret = bytearray()

        while True:
            if self.offset >= self.size or (size is not None and len(ret) >= size):
                break

            self.ext4.fd.seek(
                (
                    self.extents[self.offset_extents].start
                    + self.offset_extent_len
                )
                * self.ext4.block_size
                + self.offset % self.ext4.block_size
            )

            block_can_read = self.ext4.block_size - (self.offset % self.ext4.block_size)
            if (self.size - self.offset) < block_can_read:
                block_can_read = self.size - self.offset

            read_size = min(block_can_read, size)
            ret.extend(self.ext4.fd.read(read_size))

            old_block_nr = self.offset // self.ext4.block_size
            self.offset += read_size

            if self.offset // self.ext4.block_size > old_block_nr:
                self.offset_extent_len += (
                    self.offset // self.ext4.block_size - old_block_nr
                )

            if (
                self.offset_extent_len
                >= self.extents[self.offset_extents].len
            ):
                self.offset_extents += 1
                self.offset_extent_len = 0

            if self.offset_extents >= len(self.extents):
                self.offset = self.size
                break

        return bytes(ret)

    def listdir(self):
        """ref: htree_dirblock_to_tree"""
        assert self.mode.isdir(), "mode is %s" % self.mode

        ret = []

        for _ in range(0, self.size.value, self.ext4.block_size):
            bytes = self.read(self.ext4.block_size)
            offset = 0

            while offset < self.ext4.block_size:
                f = InodeFile(
                    bytes[offset : offset + InodeFile.length], load_name=False
                )
                if offset + f.rec_len > self.ext4.block_size:
                    break
                if f.name_len < 1:
                    break

                nbase = offset + InodeFile.length

                f.set_name(bytes[nbase : nbase + f.rec_len - InodeFile.length])

                offset += f.rec_len
                ret.append(f)

        return ret

    def _getxattr_from_bytes(self, bytes, frome_node=False):
        offset = 4 if frome_node is True else 0x20

        while offset < len(bytes) and data.le32(bytes[offset : offset + 4]) != 0:
            xattr = InodeXttr(bytes[offset:], load_value=frome_node)

            if frome_node is False:
                xattr.set_value(bytes)
            offset += (xattr.name_len + InodeXttr.length + 3) & ~3

            yield xattr

    def getxattr(self, name=None):
        ret = {}

        # ref: ext4_xattr_ibody_get
        for xattr in self._getxattr_from_bytes(self.extra_data, True):
            if name is not None and name == xattr.name:
                return xattr.valueattr

            ret[xattr.name] = xattr.value

        # ref: ext4_xattr_block_get
        if self.file_acl > 0:
            self.ext4.fd.seek(self.file_acl * self.ext4.block_size)
            bytes = self.ext4.fd.read(self.ext4.block_size)

            for xattr in self._getxattr_from_bytes(bytes, False):
                if name is not None and name == xattr.name:
                    return xattr.valueattr

            ret[xattr.name] = xattr.value

        if name is not None:
            return None
        else:
            return ret

    def seek(self, offset):
        bnr = offset // self.ext4.block_size

        offset_extents = None
        bcnt = 0
        for i in range(len(self.extents)):
            if bnr < self.extents[i].len:
                offset_extents = i
                break
            bcnt += self.extents[i].len

        if offset_extents is None:
            raise ValueError("offset %d is out of range" % offset)

        self.offset_extents = offset_extents
        self.offset_extent_len = bnr - bcnt
        self.offset = offset % self.ext4.block_size

    def readlink(self):
        assert self.mode.islnk()
 
        if self.size <= self.blocks.length:
            return self.block.link
        else:
            return self.read().decode()

class Ext4:
    def __init__(self, dev, sb: int = None) -> None:
        self.fd = open(dev, "rb")

        """
        ref: https://www.kernel.org/doc/html/v4.19/filesystems/ext4/ondisk/index.html#layout

        For the special case of block group 0, the first 1024 bytes are unused, to allow for the installation of x86 boot sectors and other oddities. 
        The superblock will start at offset 1024 bytes, whichever block that happens to be (usually 0). 
        However, if for some reason the block size = 1024, then block 0 is marked in use and the superblock goes in block 1.
        For all other block groups, there is no padding.

        对于块组 0 的特殊情况，前 1024 个字节未使用，以允许安装 x86 引导扇区和其他奇怪的东西。
        超级块将从偏移量 1024 字节开始，无论是哪个块（通常是 0）。
        但是，如果由于某种原因块大小 = 1024，则块 0 被标记为正在使用，并且超级块位于块 1 中。
        对于所有其他块组，没有填充。
        """
        self.sb_seek = sb if sb is not None else 1024

        self.logical_sb_block = 0
        self.fd.seek(self.sb_seek)  # EXT4_MIN_BLOCK_SIZE 1024

        self.sb = SuperBlock(self.fd.read(SuperBlock.length))

        self.block_size = 1024 * (2**self.sb.log_block_size.value)

        self.fd.seek(self.sb.mmp_block * self.block_size)
        self.mmp = data.MMP(self.fd.read(data.MMP.length))

        self.groups_count = int(
            (
                self.sb.blocks_count
                - self.sb.first_data_block
                + self.sb.blocks_per_group
                - 1
            )
            / self.sb.blocks_per_group
        )

        self.groups = []
        self.fd.seek(self.group_loc() * self.block_size)
        for _ in range(self.groups_count):
            self.groups.append(Group(self.fd.read(Group.length)))

    def group_loc(self):
        """ref: ext4_fsblk_t descriptor_loc, not inaccurate"""
        if self.block_size <= 1024:  # super block at 1
            return 2
        else:
            return 1

    def inode(self, nr):
        nr -= 1  # inode start at 1
        assert nr <= self.sb.inodes_count, "inode number out of range"

        group_nr = nr // self.sb.inodes_per_group
        group = self.groups[group_nr]

        inode_bitmap = group.inode_bitmap
        offset = nr % self.sb.inodes_per_group

        self.fd.seek(self.block_size * inode_bitmap + offset // 8)
        inode_bit = self.fd.read(1)[0]
        assert (inode_bit & (1 << (offset % 8))) != 0, "inode not allocated"

        inode_table = self.block_size * group.inode_table
        self.fd.seek(inode_table + offset * self.sb.inode_size)

        return Inode(self.fd.read(self.sb.inode_size.value), self)